"""
Comprehensive integration example for the event-driven continual learning framework.

This example demonstrates:
1. How to set up the complete event system with ContinualModule
2. How to use event-driven strategies (EWC and Experience Replay)
3. How to monitor and observe the training process through events
4. How to extend the system with custom event handlers

Run this example to see the complete event-driven CL pipeline in action.
"""

import torch
import torch.nn as nn
import pytorch_lightning as pl
from torch.utils.data import Dataset, DataLoader, TensorDataset
from typing import Dict, List, Any
import logging

# Import the event system components
from continuallearning.events import (
    # Core event system
    EventType,
    EventFactory,
    global_dispatcher,
    AdvancedEventDispatcher,
    # Event handlers and components
    FunctionalEventHandler,
    EventSystemCallback,
    # Event-driven strategies
    EventDrivenEWC,
    EventDrivenExperienceReplay,
)

# Import other framework components
from continuallearning.module import ContinualModule
from continuallearning.data import ContinualDataModule
from continuallearning.callbacks import TaskBoundaryCallback, TaskTrainingCallback
from continuallearning.config.typed_config import OptimizerConfig


# Configure logging for detailed output
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)


class SimpleClassificationModel(nn.Module):
    """Simple model for demonstration."""

    def __init__(
        self, input_size: int = 784, hidden_size: int = 256, num_classes: int = 10
    ):
        super().__init__()
        self.classifier = nn.Sequential(
            nn.Linear(input_size, hidden_size),
            nn.ReLU(),
            nn.Dropout(0.2),
            nn.Linear(hidden_size, hidden_size),
            nn.ReLU(),
            nn.Dropout(0.2),
            nn.Linear(hidden_size, num_classes),
        )

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return self.classifier(x.view(x.size(0), -1))


def create_synthetic_dataset(
    num_samples: int = 1000, input_size: int = 784, num_classes: int = 10
) -> TensorDataset:
    """Create a synthetic dataset for demonstration."""
    x = torch.randn(num_samples, input_size)
    y = torch.randint(0, num_classes, (num_samples,))
    return TensorDataset(x, y)


def create_task_datasets(
    num_tasks: int = 3, samples_per_task: int = 1000
) -> Dict[int, Dict[str, Dataset]]:
    """Create synthetic datasets for multiple tasks."""
    task_datasets = {}

    for task_id in range(num_tasks):
        # Create train/val/test splits for each task
        train_dataset = create_synthetic_dataset(samples_per_task, num_classes=10)
        val_dataset = create_synthetic_dataset(samples_per_task // 5, num_classes=10)
        test_dataset = create_synthetic_dataset(samples_per_task // 5, num_classes=10)

        task_datasets[task_id] = {
            "train": train_dataset,
            "val": val_dataset,
            "test": test_dataset,
        }

    logger.info(
        f"Created {num_tasks} synthetic tasks with {samples_per_task} samples each"
    )
    return task_datasets


class TrainingMonitoringHandler(FunctionalEventHandler):
    """Custom event handler to monitor training progress."""

    def __init__(self):
        super().__init__(
            event_types=[
                EventType.TASK_STARTED,
                EventType.TASK_COMPLETED,
                EventType.LOSS_COMPUTATION_COMPLETED,
                EventType.METRICS_COMPUTED,
                EventType.ERROR_OCCURRED,
            ],
            handler_func=self.handle_event,
            name="TrainingMonitoringHandler",
        )
        self.task_losses = {}
        self.task_metrics = {}

    def handle_event(self, event) -> None:
        """Handle various training events for monitoring."""
        if event.event_type == EventType.TASK_STARTED:
            logger.info(f"🚀 Task {event.task_id} started: {event.task_name}")

        elif event.event_type == EventType.TASK_COMPLETED:
            logger.info(f"✅ Task {event.task_id} completed successfully")

        elif event.event_type == EventType.LOSS_COMPUTATION_COMPLETED:
            if event.response_data and "loss_dict" in event.response_data:
                loss_dict = event.response_data["loss_dict"]
                task_id = event.task_id or 0

                if task_id not in self.task_losses:
                    self.task_losses[task_id] = []

                # Store loss information
                loss_info = {
                    k: v.item() if torch.is_tensor(v) else v
                    for k, v in loss_dict.items()
                }
                self.task_losses[task_id].append(loss_info)

                # Log detailed loss information
                logger.info(f"📊 Task {task_id} Loss: {loss_info}")

        elif event.event_type == EventType.METRICS_COMPUTED:
            if event.metadata and "metrics" in event.metadata:
                metrics = event.metadata["metrics"]
                task_id = event.task_id or 0

                if task_id not in self.task_metrics:
                    self.task_metrics[task_id] = []

                self.task_metrics[task_id].append(metrics)
                logger.info(f"📈 Task {task_id} Metrics: {metrics}")

        elif event.event_type == EventType.ERROR_OCCURRED:
            logger.error(f"❌ Error in {event.source}: {event.error_message}")

    def get_summary(self) -> Dict[str, Any]:
        """Get a summary of collected metrics."""
        return {
            "task_losses": self.task_losses,
            "task_metrics": self.task_metrics,
            "total_tasks_processed": len(self.task_losses),
        }


class MemoryUsageHandler(FunctionalEventHandler):
    """Handler to monitor memory usage in experience replay."""

    def __init__(self):
        super().__init__(
            event_types=[
                EventType.MEMORY_UPDATE_COMPLETED,
                EventType.BATCH_PREPROCESSING_COMPLETED,
            ],
            handler_func=self.handle_event,
            name="MemoryUsageHandler",
        )

    def handle_event(self, event) -> None:
        """Monitor memory-related events."""
        if event.event_type == EventType.MEMORY_UPDATE_COMPLETED:
            if event.response_data:
                data = event.response_data
                logger.info(
                    f"🧠 Memory Updated - Task {event.task_id}: "
                    f"Added {data.get('samples_added', 0)} samples, "
                    f"Memory size: {data.get('memory_size', 0)}/{data.get('memory_capacity', 0)}"
                )

        elif event.event_type == EventType.BATCH_PREPROCESSING_COMPLETED:
            if event.response_data:
                data = event.response_data
                logger.info(
                    f"🔄 Batch Augmented - Task {event.task_id}: "
                    f"Added {data.get('replay_samples_added', 0)} replay samples, "
                    f"Final batch size: {data.get('final_batch_size', 0)}"
                )


def setup_event_monitoring():
    """Set up comprehensive event monitoring."""
    # Create custom event handlers
    training_monitor = TrainingMonitoringHandler()
    memory_monitor = MemoryUsageHandler()

    # Register handlers with global dispatcher
    global_dispatcher.register_handler(training_monitor)
    global_dispatcher.register_handler(memory_monitor)

    logger.info("✅ Event monitoring handlers registered")

    return training_monitor, memory_monitor


def create_event_driven_cl_system():
    """Create a complete event-driven continual learning system."""

    # 1. Set up event monitoring
    training_monitor, memory_monitor = setup_event_monitoring()

    # 2. Create the model
    model = SimpleClassificationModel(input_size=784, hidden_size=256, num_classes=10)

    # 3. Create event-driven strategies
    ewc_strategy = EventDrivenEWC(
        regularization_weight=5000.0,
        fisher_estimation_samples=100,  # Smaller for demo
        name="EWC_Strategy",
    )

    replay_strategy = EventDrivenExperienceReplay(
        memory_size=500,  # Smaller for demo
        replay_batch_size=16,
        name="ExperienceReplay_Strategy",
    )

    # Note: We'll use EWC as primary strategy, but both are registered with the event system

    # 4. Create ContinualModule with event integration
    optimizer_config = OptimizerConfig(type="Adam", lr=0.001, weight_decay=1e-4)

    continual_module = ContinualModule(
        model=model,
        optimizer_cfg=optimizer_config,
        cl_strategy=ewc_strategy,  # Primary strategy
        criterion=nn.CrossEntropyLoss(),
        event_dispatcher=global_dispatcher,
    )

    # 5. Create task datasets
    task_datasets = create_task_datasets(
        num_tasks=3, samples_per_task=500
    )  # Smaller for demo

    # 6. Create continual data module
    continual_datamodule = ContinualDataModule(
        task_datasets=task_datasets,
        batch_size=32,
        num_workers=0,  # Single-threaded for demo
        shuffle=True,
    )

    # 7. Create Lightning callbacks with event integration
    event_callback = EventSystemCallback(dispatcher=global_dispatcher)
    task_boundary_callback = TaskBoundaryCallback(
        verbose=True, event_dispatcher=global_dispatcher
    )
    task_training_callback = TaskTrainingCallback(
        task_boundary_callback=task_boundary_callback,
        reset_optimizer_between_tasks=False,
    )

    # 8. Create Lightning trainer with all callbacks
    trainer = pl.Trainer(
        max_epochs=2,  # Small number for demo
        callbacks=[event_callback, task_boundary_callback, task_training_callback],
        enable_progress_bar=True,
        enable_model_summary=True,
        logger=False,  # Disable wandb/tensorboard for demo
    )

    logger.info("✅ Complete event-driven CL system created")

    return {
        "trainer": trainer,
        "module": continual_module,
        "datamodule": continual_datamodule,
        "task_training_callback": task_training_callback,
        "monitors": {"training": training_monitor, "memory": memory_monitor},
        "strategies": {"ewc": ewc_strategy, "replay": replay_strategy},
    }


def run_continual_learning_example():
    """Run the complete continual learning example."""

    logger.info("🚀 Starting Event-Driven Continual Learning Example")

    # Create the complete system
    system = create_event_driven_cl_system()

    trainer = system["trainer"]
    continual_module = system["module"]
    continual_datamodule = system["datamodule"]
    task_training_callback = system["task_training_callback"]
    monitors = system["monitors"]

    # Prepare data for each task
    tasks_data = []
    for task_id in range(continual_datamodule.num_tasks):
        continual_datamodule.set_task(task_id)

        train_loader = continual_datamodule.train_dataloader()
        val_loader = continual_datamodule.val_dataloader()

        tasks_data.append({"train": train_loader, "val": val_loader})

    logger.info(f"📚 Prepared data for {len(tasks_data)} tasks")

    # Run continual learning across all tasks
    logger.info("🎯 Starting continual learning training...")

    try:
        results = task_training_callback.train_multiple_tasks(
            trainer=trainer, model=continual_module, tasks_data=tasks_data
        )

        logger.info("✅ Continual learning training completed successfully!")

        # Print summary of results
        print("\n" + "=" * 60)
        print("CONTINUAL LEARNING RESULTS SUMMARY")
        print("=" * 60)

        training_summary = monitors["training"].get_summary()
        print(f"📊 Total tasks processed: {training_summary['total_tasks_processed']}")

        for task_id, losses in training_summary["task_losses"].items():
            if losses:
                final_loss = losses[-1]
                print(f"Task {task_id} final loss: {final_loss}")

        # Show EWC strategy state
        ewc_strategy = system["strategies"]["ewc"]
        print(f"\n🧠 EWC Strategy State:")
        print(f"  - Completed tasks: {list(ewc_strategy.state.completed_tasks)}")
        print(
            f"  - Fisher information computed for: {len(ewc_strategy.state.fisher_information)} parameters"
        )
        print(
            f"  - Optimal parameters stored for: {len(ewc_strategy.state.optimal_params)} parameters"
        )

        # Show Experience Replay state
        replay_strategy = system["strategies"]["replay"]
        print(f"\n🔄 Experience Replay State:")
        print(f"  - Memory buffer size: {len(replay_strategy.memory_buffer)}")
        print(f"  - Task distribution: {replay_strategy._get_task_distribution()}")

        print("=" * 60)

        return results

    except Exception as e:
        logger.error(f"❌ Error during continual learning: {e}")
        raise


def demonstrate_event_system_features():
    """Demonstrate additional event system features."""

    logger.info("🔧 Demonstrating Event System Features")

    # Show dispatcher statistics
    stats = global_dispatcher.stats
    if stats:
        print(f"\n📈 Event Dispatcher Statistics:")
        print(f"  - Events processed: {stats.events_processed}")
        print(f"  - Events failed: {stats.events_failed}")
        print(f"  - Handlers registered: {stats.handlers_registered}")
        print(f"  - Filters registered: {stats.filters_registered}")

    # Show registered handlers
    handler_count = 0
    for event_type, handlers in global_dispatcher._handlers.items():
        if handlers:
            print(f"  - {event_type.name}: {len(handlers)} handlers")
            handler_count += len(handlers)

    print(f"  - Total handlers across all events: {handler_count}")

    # Test custom event creation and dispatch
    logger.info("\n🧪 Testing custom event creation...")

    custom_event = (
        EventFactory.logging(log_level="INFO", message="Custom event test from example")
        .source("integration_example")
        .tag("test")
        .build()
    )

    global_dispatcher.dispatch(custom_event)
    logger.info("✅ Custom event dispatched successfully")


if __name__ == "__main__":
    """Run the complete example."""

    try:
        # Run the main continual learning example
        results = run_continual_learning_example()

        # Demonstrate additional features
        demonstrate_event_system_features()

        print("\n🎉 Event-driven continual learning example completed successfully!")
        print("This example demonstrates:")
        print("  ✅ Complete event system integration")
        print("  ✅ Event-driven strategy implementations")
        print("  ✅ Custom event handlers and monitoring")
        print("  ✅ Lightning callback integration")
        print("  ✅ Multi-task continual learning pipeline")

    except Exception as e:
        logger.error(f"💥 Example failed: {e}")
        raise
